# -*- coding: utf-8 -*-
"""
HHL (cqlib version) — with optional Amplitude Amplification (AA)
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List, Tuple, Union
import numpy as np
import math

# cqlib 基础组件
from circuits.algorithm_circuit import Algorithm_Circuit as Circuit
from cqlib.simulator import StatevectorSimulator

# 相位估计（QPE）相关工具
from algorithm.quantum_phase_estimator import (
    hhl_pe_encode, hhl_pe_uncompute,
    pick_t_from_lambda_max, rough_lambda_max_via_eigs,
)

# 哈密顿演化 oracle（Pauli 分解 + Trotter/Suzuki）
from algorithm.Hamiltonian_simulator import (
    build_pauli_evolution_oracle_from_matrix,
)

# 倒数旋转模块（LookupRotation）
from algorithm.lookup_rotation import LookupRotation

# 向量态制备模块（幅度加载）
from algorithm.quantum_state_prep import prepare_state_from_vector

# Grover / AA 组件（你已有的 cqlib 版本）
from algorithm.grover_operator import grover_operator


# ---------------------------
# 结果对象（容器）
# ---------------------------
@dataclass
class HHLResult:
    solution: Optional[np.ndarray] = None
    output: Optional[np.ndarray] = None
    probability_result: Optional[float] = None
    matrix: Optional[np.ndarray] = None
    vector: Optional[np.ndarray] = None
    circuit_info: Dict[str, Any] = field(default_factory=dict)
    aa_info: Dict[str, Any] = field(default_factory=dict)  # 新增：AA 统计信息


# ---------------------------
# 工具：统一态矢量为一维复 ndarray
# ---------------------------
def _coerce_statevector_to_array(sv_like, total_n: int) -> np.ndarray:
    N = 1 << int(total_n)
    if isinstance(sv_like, dict):
        arr = np.zeros(N, dtype=complex)
        for k, v in sv_like.items():
            i = int(k, 2) if isinstance(k, str) else int(k)
            arr[i] = complex(v)
        return arr
    return np.asarray(sv_like, dtype=complex).reshape(N)


# ---------------------------
# HHL 主类（加入 AA 配置）
# ---------------------------
class HHL:
    def __init__(
        self,
        matrix: np.ndarray,
        vector: np.ndarray,
        m_phase: int = 4,
        trotter_slices: int = 2,
        truncate_powerdim: bool = False,
        truncate_hermitian: bool = False,
        orig_size: Optional[int] = None,
        # === AA 参数（新增） ===
        aa_steps: int = 0,
        aa_reflection: Union[str, List[int]] = "anc_only",  # "anc_only" | "all" | [qubit indices]
        aa_insert_barriers: bool = False,
    ):
        """
        aa_steps:
            0 表示关闭幅度放大；>0 表示重复 Q=A S0 A† Sf 的次数。
        aa_reflection:
            - "anc_only": 仅对 anc 做 S0（工程上稳健、资源开销极小，推荐默认）
            - "all": 对全部比特做 S0（需要大规模多控-Z 与充足 ancilla）
            - [list]: 自定义 S0 作用的比特集合
        aa_insert_barriers:
            可视化/调试用，在 Grover 关键阶段加入 barrier。
        """
        # ---------- 原有初始化 ----------
        A = np.array(matrix, dtype=float)
        b = np.array(vector, dtype=float).reshape(-1)
        if A.shape[0] != A.shape[1]:
            raise ValueError("matrix must be square")
        if A.shape[0] != b.size:
            raise ValueError("vector dimension mismatch")

        self._orig_A = A.copy()
        self._orig_b = b.copy()
        self._truncate_powerdim = truncate_powerdim
        self._truncate_hermitian = truncate_hermitian
        self._orig_size = orig_size or b.size

        A2, b2 = A, b
        if truncate_powerdim and (math.log2(A.shape[0]) % 1 != 0):
            A2, b2 = self.expand_to_powerdim(A2, b2)
        if truncate_hermitian and not np.allclose(A2, A2.conj().T):
            A2, b2 = self.expand_to_hermitian(A2, b2)

        n_qubits_work = int(round(math.log2(A2.shape[0])))
        if 2 ** n_qubits_work != A2.shape[0]:
            raise ValueError("matrix dimension must be 2^n after resize")
        if not np.allclose(A2, A2.conj().T):
            raise ValueError("matrix must be Hermitian after resize")

        self.A = A2
        self.b = b2 / np.linalg.norm(b2)
        self.n_work = n_qubits_work
        self.m_phase = int(m_phase)
        self.trotter_slices = int(trotter_slices)

        # 缓存
        self._qc: Optional[Circuit] = None
        self._reg: Optional[Dict[str, List[int]]] = None
        self._t: Optional[float] = None

        # AA 配置
        self.aa_steps = int(aa_steps)
        self.aa_reflection = aa_reflection
        self.aa_insert_barriers = bool(aa_insert_barriers)

    # ---------- 扩维工具 ----------
    @staticmethod
    def expand_to_powerdim(matrix: np.ndarray, vector: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        dim = matrix.shape[0]
        next_n = int(np.ceil(np.log2(dim)))
        newA = np.eye(2 ** next_n, dtype=float)
        newA[:dim, :dim] = matrix
        newb = np.zeros(2 ** next_n, dtype=float)
        newb[:vector.size] = vector
        return newA, newb

    @staticmethod
    def expand_to_hermitian(matrix: np.ndarray, vector: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        d = matrix.shape[0]
        A = np.zeros((2 * d, 2 * d), dtype=float)
        A[:d, d:] = matrix.conj().T
        A[d:, :d] = matrix
        b = np.zeros(2 * d, dtype=float)
        b[:d] = vector.conj()
        b[d:] = vector
        return A, b

    # ---------- 构造（可选 AA） ----------
    def construct_circuit(self, measurement: bool = False) -> Circuit:
        """
        当 aa_steps==0：构造单次 HHL 子程序（与旧版一致）。
        当 aa_steps>0 ：构造 A · (Q)^aa_steps，其中 Q = A S0 A† Sf。
        """
        A = self.A
        b = self.b
        n_work = self.n_work
        m_phase = self.m_phase

        # 选择 t
        lam_max = rough_lambda_max_via_eigs(A)
        t = pick_t_from_lambda_max(lam_max, target_phase=0.25, safety=0.95)

        # e^{-iAt} oracle
        evo = build_pauli_evolution_oracle_from_matrix(A, method='suzuki2', num_slices=self.trotter_slices)

        # LookupRotation 预估需要的额外位
        aux_needed_lookup = max(0, m_phase - 1)
        extra_lookup = 3 + aux_needed_lookup

        total_qubits = m_phase + n_work + extra_lookup
        qc_A = Circuit(total_qubits)

        # 寄存器划分
        work = list(range(0, n_work))
        phase = list(range(n_work, n_work + m_phase))

        # |b> 制备
        prep, _, _, _ = prepare_state_from_vector(list(b))
        qc_A = qc_A.compose(prep, work)

        # QPE 编码
        hhl_pe_encode(qc_A, phase_qubits=phase, work_qubits=work, evolution_oracle=evo, t=t)

        # 倒数旋转
        lam_vals = np.linalg.eigvalsh(A)
        lam_min = float(np.min(np.abs(lam_vals)))
        self._lookup = LookupRotation(
            pat_length=None, subpat_length=None,
            negative_evals=False, evo_time=t, lambda_min=lam_min,
        )
        qc_A = self._lookup.construct_circuit(mode="circuit", inreg=phase, circuit=qc_A)

        # QPE 回滚
        hhl_pe_uncompute(qc_A, phase_qubits=phase, work_qubits=work, evolution_oracle=evo, t=t)

        # anc / aux
        anc = int(self._lookup._anc)  # type: ignore[attr-defined]
        aux_lookup = list(getattr(self._lookup, "_aux_pool", []))  # type: ignore[attr-defined]

        # measurement 预留
        if measurement:
            pass

        # ====== 若不开 AA：直接缓存并返回 ======
        if self.aa_steps <= 0:
            self._qc = qc_A
            self._reg = {"work": work, "phase": phase, "anc": [anc], "aux": aux_lookup}
            return qc_A

        # ====== 构造 Oracle：标记 anc=1 的好态 ======
        oracle = Circuit(total_qubits)
        oracle.z(anc)  # S_f：对 anc=1 加相位 π

        # ====== 选择反射位集合 S0 的比特 ======
        if isinstance(self.aa_reflection, list):
            reflection_qubits = list(self.aa_reflection)
        elif isinstance(self.aa_reflection, str) and self.aa_reflection.lower() == "all":
            reflection_qubits = list(range(total_qubits))
        else:
            # 默认仅对 anc 做零反射（X-Z-X）
            reflection_qubits = [anc]

        # ====== Grover 扩散算子 Q ======
        # 需要 A 的逆；若库未实现 inverse() 请补齐 Algorithm_Circuit.inverse()
        if not hasattr(qc_A, "inverse"):
            raise NotImplementedError(
                "需要 Algorithm_Circuit.inverse() 用于构造 Q = A S0 A† Sf；"
                "请为 cqlib 电路实现 inverse() 或将 aa_steps 设为 0 关闭 AA。"
            )
        grover_Q = grover_operator(
            oracle=oracle,
            state_preparation=qc_A,
            reflection_qubits=reflection_qubits,
            insert_barriers=self.aa_insert_barriers,
            name="HHL-AA-Q"
        )

        # ====== 拼接最终电路：A · Q^k ======
        qc_full = Circuit(total_qubits)
        qc_full = qc_full.compose(qc_A)
        for _ in range(self.aa_steps):
            qc_full = qc_full.compose(grover_Q)

        self._qc = qc_full
        self._reg = {"work": work, "phase": phase, "anc": [anc], "aux": aux_lookup}
        return qc_full

    # ---------- 从 statevector 提取解（条件于 anc=1） ----------
    @staticmethod
    def _extract_solution_from_statevector(
        sv_like, n_phase: int, n_work: int, anc_index: int, total_n: int
    ) -> np.ndarray:
        sv = np.asarray(sv_like, dtype=complex).reshape(-1)
        x_vec = np.zeros(1 << n_work, dtype=complex)
        for idx in range(1 << total_n):
            if ((idx >> anc_index) & 1) != 1:
                continue
            w_small = 0
            for j in range(n_work):
                w_small |= (((idx >> j) & 1) << j)
            x_vec[w_small] += sv[idx]
        nrm = np.linalg.norm(x_vec)
        return x_vec / nrm if nrm > 0 else x_vec

    # ---------- 运行（statevector 仿真） ----------
    def run(self, statevector: bool = True) -> HHLResult:
        if self._qc is None:
            self.construct_circuit(measurement=False)

        qc = self._qc
        reg = self._reg or {}
        phase = reg.get("phase", [])
        work  = reg.get("work", [])
        anc   = reg.get("anc", [])[0]

        sim = StatevectorSimulator(qc)
        sv_raw = sim.statevector()
        total_n = qc.num_qubits
        sv_arr = _coerce_statevector_to_array(sv_like=sv_raw, total_n=total_n)

        # anc=1 成功概率
        prob_anc1 = 0.0
        for idx in range(1 << total_n):
            if ((idx >> anc) & 1) == 1:
                prob_anc1 += float(abs(sv_arr[idx]) ** 2)

        # 量子方向解（条件于 anc=1）
        x_quantum = self._extract_solution_from_statevector(
            sv_like=sv_arr, n_phase=len(phase), n_work=len(work), anc_index=anc, total_n=total_n
        )

        # 与经典方向对齐（去全局相位）
        x_classic_dir = np.linalg.solve(self.A, self.b)
        x_classic_dir = x_classic_dir / np.linalg.norm(x_classic_dir)
        phase_corr = np.vdot(x_classic_dir, x_quantum)
        if abs(phase_corr) > 0:
            x_quantum *= phase_corr / abs(phase_corr)

        # 取实方向并归一化（实对称 A 情况）
        output = x_quantum.copy()
        x_aligned = output * np.exp(-1j * np.angle(output[0])) if output.size else output
        x_dir = np.real(x_aligned)
        if np.linalg.norm(x_dir) == 0 and output.size:
            x_dir = np.real(output)
        if np.linalg.norm(x_dir) > 0:
            x_dir = x_dir / np.linalg.norm(x_dir)

        # 恢复到“经典尺度”
        denom = np.linalg.norm(self.A @ x_dir) if x_dir.size else 0.0
        alpha = (np.linalg.norm(self.b) / denom) if denom > 0 else 0.0
        x_solution_full = alpha * x_dir
        output_full = output

        # 若做了扩维/厄米化，则截回原尺寸
        x_solution = x_solution_full
        output = output_full
        if self._truncate_hermitian:
            x_solution = x_solution[: self._orig_size]
            output = output[: self._orig_size]
        if self._truncate_powerdim:
            x_solution = x_solution[: self._orig_size]
            output = output[: self._orig_size]

        circuit_info = {
            "num_qubits": qc.num_qubits,
            "num_instructions": getattr(qc, "num_instructions", None)
        }

        aa_info = {
            "aa_steps": self.aa_steps,
            "aa_reflection": self.aa_reflection if isinstance(self.aa_reflection, str) else "custom",
            "aa_insert_barriers": self.aa_insert_barriers,
        }

        return HHLResult(
            solution=x_solution,
            output=output,
            probability_result=float(prob_anc1),
            matrix=self._orig_A,
            vector=self._orig_b,
            circuit_info=circuit_info,
            aa_info=aa_info,
        )


# ------------- 简易自测 -------------
if __name__ == "__main__":
    A = np.array([[1, 1/3], [1/3, 1]], dtype=float)
    b = np.array([1.0, 0], dtype=float)

    # 无 AA
    hhl0 = HHL(A, b, m_phase=8, trotter_slices=2, aa_steps=0)
    hhl0.construct_circuit()
    r0 = hhl0.run()

    # 开 AA（仅对 anc 反射，开销极小）
    hhla = HHL(A, b, m_phase=8, trotter_slices=2, aa_steps=2, aa_reflection="anc_only")
    hhla.construct_circuit()
    ra = hhla.run()

    x_true = np.linalg.solve(A, b)
    rel0 = np.linalg.norm(r0.solution - x_true) / np.linalg.norm(x_true)
    rela = np.linalg.norm(ra.solution - x_true) / np.linalg.norm(x_true)

    print("=== HHL (baseline) ===")
    print("prob anc=1 =", r0.probability_result, "rel.err =", f"{rel0:.6%}")
    print("=== HHL + AA (k=2, anc_only) ===")
    print("prob anc=1 =", ra.probability_result, "rel.err =", f"{rela:.6%}")
